package cn.kgc.config;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * 使用Zookeeper原生API实现分布式锁
 *
 * @author YC
 */
public class ZookeeperLock {

    /**
     * 倒计时闩锁，初始值1
     */
    private final CountDownLatch countDownLatch = new CountDownLatch(1);

    /**
     * zookeeper连接
     */
    private final ZooKeeper zooKeeper;
    /**
     * 锁的超时时间
     */
    private final int lockTimeOut;
    /**
     * 锁名
     */
    private final String lockName;
    /**
     * 锁命名空间
     */
    private final String lockSpace;
    /**
     * 当前锁的路径
     */
    private String selfLockPath;

    /**
     * 构造函数-创建分布式锁
     */
    public ZookeeperLock(ZooKeeper zooKeeper, int lockTimeOut, String lockName, String lockSpace) {
        this.zooKeeper = zooKeeper;
        this.lockTimeOut = lockTimeOut;
        this.lockName = lockName;
        this.lockSpace = lockSpace;
        // 创建分布式锁的根命名空间
        try {
            Stat stat = zooKeeper.exists(lockSpace, false);
            if (stat == null) {
                // 如果没有创建根节点
                zooKeeper.create(lockSpace, "".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (Exception e) {
            throw new RuntimeException("创建分布式锁根命名空间失败");
        }
    }

    /**
     * 获取锁
     */
    public void lock() throws Exception {
        // 创建锁节点
        selfLockPath = zooKeeper.create(
                lockSpace + "/" + lockName,
                "".getBytes(),
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT_SEQUENTIAL
        );
        // 记录当前时间
        long startTime = System.currentTimeMillis();
        // 获得当前锁的节点名称   /lockSpace/lockName0000000001 => lockName0000000001
        String lockNode = selfLockPath.substring(lockSpace.length() + 1);
        // 尝试获得锁，在循环中尝试获得锁，获得跳出循环
        while (true) {
            // 当前时间 - 开始时间 > 超时时间   则超时
            if (System.currentTimeMillis() - startTime > lockTimeOut) {
                throw new RuntimeException("获取锁超时");
            }
            // 判断是否获得了锁（锁命名空间下是否有子节点）
            List<String> children = zooKeeper.getChildren(lockSpace, false);
            if (children.size() == 1) {
                return;
            }
            // 如果有多个字节点，获取自己创建的节点的顺序
            List<String> sort = children.stream()
                    .map(path -> path.substring(lockSpace.length() + 1))
                    .sorted()
                    .collect(Collectors.toList());
            // 监听排序在前面的节点
            int preNodeIndex = sort.indexOf(lockNode) - 1;
            String preNode = sort.get(preNodeIndex);
            Stat stat = zooKeeper.exists(lockSpace + "/" + preNode, new Watcher() {
                @Override
                public void process(WatchedEvent watchedEvent) {
                    if (watchedEvent.getType() == Event.EventType.NodeDeleted) {
                        // 前一个节点释放锁，重新尝试获得锁
                        countDownLatch.countDown();
                    }
                }
            });
            if (stat != null) {
                // 距离锁超时的剩余时间
                int remainTime = lockTimeOut - (int) (System.currentTimeMillis() - startTime);
                if (remainTime < 0) {
                    throw new RuntimeException("获取锁超时");
                }
                // 在锁的超时时间内，等待前一个节点释放锁
                boolean await = countDownLatch.await(remainTime, TimeUnit.MILLISECONDS);
            }
            // 继续下一个循环
        }
    }

    /**
     * 释放锁
     */
    public void unLock() {
        if (selfLockPath != null) {
            try {
                // 删除根命名空间节点，删除不考虑版本
                zooKeeper.delete(selfLockPath, -1);
                selfLockPath = null;
            } catch (Exception e) {
                throw new RuntimeException("释放锁失败");
            }
        }
    }

}
